
import torch
import numpy as np

def myIOUnp(pic0xy, pic1xy, useMinIou=False):
    """
    :param pic0xy: 第一张图片左上角坐标和右下角坐标，输入格式必须为元组，里面存放的格式为x0,y0,x1,y1
    :param pic1xy: 第二张图片左上角坐标和右下角坐标
    :return:
    """
    # 左上角x,y，要分别取，反正就是取左上角坐标中大的那个
    minx = np.maximum(pic0xy[:, 0], pic1xy[:, 0])
    miny = np.maximum(pic0xy[:, 1], pic1xy[:, 1])
    maxx = np.minimum(pic0xy[:, 2], pic1xy[:, 2])
    maxy = np.minimum(pic0xy[:, 3], pic1xy[:, 3])
    # 算交集面积 长宽
    awidth = maxx - minx
    aheight = maxy - miny
    awidth[awidth <= 0] = 0
    aheight[aheight <= 0] = 0
    area1 = (awidth) * (aheight)

    # 是否取较小值，并集除以较小值
    if useMinIou:
        # 计算交集 除以较小面积，大框套下框的情况
        pt0area = (pic0xy[:, 2] - pic0xy[:, 0]) * (pic0xy[:, 3] - pic0xy[:, 1])
        pt1area = (pic1xy[:, 2] - pic1xy[:, 0]) * (pic1xy[:, 3] - pic1xy[:, 1])
        minarea = np.minimum(pt0area, pt1area)
        return area1 / minarea
    else:
        # 算并集面积 各自面积相加-交集
        area2 = (pic0xy[:, 2] - pic0xy[:, 0]) * (pic0xy[:, 3] - pic0xy[:, 1]) + (pic1xy[:, 2] - pic1xy[:, 0]) * (
                    pic1xy[:, 3] - pic1xy[:, 1]) - area1
        return area1 / area2

def myIOU(pic0xy, pic1xy, useMinIou=False):
    """
    :param pic0xy: 第一张图片左上角坐标和右下角坐标，输入格式必须为元组，里面存放的格式为x0,y0,x1,y1
    :param pic1xy: 第二张图片左上角坐标和右下角坐标
    :return:
    """
    # 左上角x,y，要分别取，反正就是取左上角坐标中大的那个
    minx = torch.max(pic0xy[:, 0], pic1xy[:, 0])
    miny = torch.max(pic0xy[:, 1], pic1xy[:, 1])
    maxx = torch.min(pic0xy[:, 2], pic1xy[:, 2])
    maxy = torch.min(pic0xy[:, 3], pic1xy[:, 3])
    # 算交集面积 长宽
    awidth = maxx - minx
    aheight = maxy - miny
    awidth[awidth <= 0] = 0
    aheight[aheight <= 0] = 0
    area1 = (awidth) * (aheight)

    # 是否取较小值，并集除以较小值
    if useMinIou:
        # 计算交集 除以较小面积，大框套下框的情况
        pt0area = (pic0xy[:, 2] - pic0xy[:, 0]) * (pic0xy[:, 3] - pic0xy[:, 1])
        pt1area = (pic1xy[:, 2] - pic1xy[:, 0]) * (pic1xy[:, 3] - pic1xy[:, 1])
        minarea = torch.min(pt0area, pt1area)
        return area1 / minarea
    else:
        # 算并集面积 各自面积相加-交集
        area2 = (pic0xy[:, 2] - pic0xy[:, 0]) * (pic0xy[:, 3] - pic0xy[:, 1]) + (pic1xy[:, 2] - pic1xy[:, 0]) * (
                    pic1xy[:, 3] - pic1xy[:, 1]) - area1
        return area1 / area2
# 需要输入形状为二维
# 传入的数据为 置信度，左上角坐标，右上角坐标
def myNMS(netout, iouthre=0.6, useMinIou=False):
    # 1.置信度筛选

    newArray = netout
    # 排序# 排序 从小到大  indices是原数组索引 用了数组索引  排序算一次就可以了
    data, indices = torch.sort(newArray, dim=0)
    newArray = newArray[indices[:, 0]]

    if newArray.numel() == torch.Size([]):
        print("this tensor is empty")
        # print(netout[:, 0] >= coef)
        return []
    target = []
    i = 0
    while True:
        # 2.找出最大置信度的框
        # 3.计算其余框和它的IOU 保留>0.6的框
        # 4.终止条件为找出所有的框
        i += 1
        if newArray.shape[0] == 1:
            target.append(newArray[0].tolist())
            break
        # 增加目标数组
        try:
            cmpitem = newArray[-1]
            target.append(cmpitem.tolist())
            newArray = newArray[:-1]
        except:
            break
        # 计算IOU
        iou = myIOU(newArray[:, 1:5], cmpitem[1:5].reshape(-1, 4), useMinIou)
        # 去除IOU大于0.6的部分
        newArray = newArray[iou <= iouthre]
    return target

def myNMSnp(netout, iouthre=0.6, useMinIou=False):
    # 1.置信度筛选

    newArray = netout
    # 排序# 排序 从小到大  indices是原数组索引 用了数组索引  排序算一次就可以了
    indices = np.argsort(newArray, axis=0)
    newArray = newArray[indices[:, 0]]

    if newArray.shape[0] == 0:
        # print("this tensor is empty")
        # print(netout[:, 0] >= coef)
        return []
    target = []
    i = 0
    while True:
        # print(newArray.shape)
        # 2.找出最大置信度的框
        # 3.计算其余框和它的IOU 保留>0.6的框
        # 4.终止条件为找出所有的框
        i += 1
        if newArray.shape[0] == 1:
            target.append(newArray[0])
            break
        # 增加目标数组
        try:
            cmpitem = newArray[-1]
            target.append(cmpitem)
            newArray = newArray[:-1]
        except:
            break
        # 计算IOU
        iou = myIOUnp(newArray[:, 1:5], cmpitem[1:5].reshape(-1, 4), useMinIou)
        # 去除IOU大于0.6的部分
        newArray = newArray[iou <= iouthre]
    return np.vstack(target)
if __name__ == '__main__':
    b = torch.tensor([[0.95, 6, 10, 20, 20], [0.8, 30, 10, 50, 20], [0.9, 0, 10, 50, 20], [0.89, 0, 9, 50, 20]],
                     dtype=torch.float32)
    a = np.array([[0.95, 6, 10, 20, 20], [0.8, 30, 10, 50, 20], [0.9, 0, 10, 50, 20], [0.89, 0, 9, 50, 20]]
                )
    print(myNMSnp(a))
    print(myNMS(b))
